TL;DR

This type-cleaning PR needed a couple of changes I'd like to get feedback on (especially from @Pouyanpi ).

1. We import the langchain_nvidia_ai_endpoints package in nemoguardrails/llm/providers/_langchain_nvidia_ai_endpoints_patch.py. To type-check against this we need to install the langchain-nvidia-ai-endpoints package. I did this by modifying pr-tests.yml and test-coverage-report.yml to add the --all-extras option to poetry install. Any concerns with adding in this amount of extra packages?
2. We have code importing tritonclient in nemoguardrails/llm/providers/trtllm/client.py. The latest version of Python supported by tritonclient is 3.8 which we deprecated on the last release. The code isn't referred to anywhere else. Should we deprecate this? I added it to the pyright exclude list while we decide what to do.

Type-Safety Fixes Summary

This report summarizes the code changes made to improve type safety. The fixes have been categorized into high, medium, and low-risk buckets based on their potential to cause regressions or unintended behavior.

🔴 High-Risk Changes

This category includes changes that alter a function's core contract, making it a potentially breaking change for consumers of the API.

Stricter Function Argument Type

This fix tightens a function's API by changing an optional argument to a required one. This is high-risk because any calling code that previously passed None will now fail.

• File: nemoguardrails/llm/models/initializer.py
• Line: 27
• Original Error: model_name was Optional[str], but the function's logic implicitly required a string value, creating a risk of a None value causing a downstream crash.
• Fix:

  def init_llm_model(
      model_name: str,
      provider_name: str,
      mode: Literal["chat", "text"],
      kwargs: Dict[str, Any],
  ) -> Union[BaseLLM, BaseChatModel]:

• Explanation: The type hint for model_name was changed from Optional[str] to str, making it a required argument. This enforces the contract at the type-checking level and makes the function's expectations explicit.
• Assumptions: It's assumed that model_name should never be None and that all call sites can provide a valid string. This assumption is supported by related changes in nemoguardrails/rails/llm/llmrails.py which now ensure a valid model name is always passed.
• Alternative Fixes: An alternative would be to keep the type as Optional[str] and raise a ValueError inside the function if model_name is None. The implemented fix is better because it fails earlier (at static analysis or integration time) rather than at runtime.

🟡 Medium-Risk Changes

These changes correct significant errors or modify function signatures in a way that is likely correct but could affect downstream code that depended on the previous (incorrect) behavior.

1. Corrected Function Return Type Annotation

The function's implementation returned a different type than what its signature declared.

• File: nemoguardrails/llm/filters.py
• Line: 278
• Original Error: The function was annotated to return a str but was actually returning a List[dict].
• Fix:

  def to_chat_messages(events: List[dict]) -> List[dict]:

• Explanation: The return type annotation was corrected from str to List[dict] to match the actual output of the function. This resolves the inconsistency and provides the correct type information to static analyzers and developers.
• Assumptions: Assumes that all callers of this function were already implicitly handling a List[dict] and that no code was relying on the incorrect str annotation.
• Alternative Fixes: Another option would be to modify the function to serialize the list of dictionaries into a string (e.g., JSON) before returning. However, correcting the annotation is the superior choice as it aligns the signature with the function's clear purpose, which is to transform events into a message list.

2. Modified Function Signature to Match Superclass

The method signature in a subclass did not match the signature of the method it was overriding in the parent class.

• File: nemoguardrails/llm/helpers.py
• Lines: 77-80 and 88-91
• Original Error: The _call and _acall methods in the WrapperLLM class were missing the **kwargs parameter present in the base LLM class from LangChain, causing a signature mismatch error (Liskov Substitution Principle violation).
• Fix:

  # For _call
  def _call(
      self,
      prompt: str,
      stop: Optional[List[str]] = None,
      run_manager: Optional[CallbackManagerForLLMRun] = None,
      **kwargs,
  ) -> str:
      self._modify_instance_kwargs()
      return llm_instance._call(prompt, stop, run_manager, **kwargs)
  
  # A similar change was made for _acall

• Explanation: **kwargs was added to the method signatures and passed through to the wrapped llm_instance call. This makes the wrapper compliant with the base class interface, allowing it to accept and forward any additional keyword arguments.
• Assumptions: It is assumed that the wrapped llm_instance also accepts **kwargs in its _call and _acall methods, which is a safe assumption for LangChain LLM objects.
• Alternative Fixes: There are no reasonable alternatives. The signature of an overridden method must be compatible with its superclass.

3. Added Control Flow for Optional Values in Core Logic

A loop condition involved comparing a value with a potentially None attribute, which would cause a TypeError at runtime.

• File: nemoguardrails/llm/taskmanager.py
• Lines: 268-269 and 294-296
• Original Error: The code compared len(task_prompt) with prompt.max_length inside a while loop without checking if prompt.max_length was None.
• Fix:

  while (
      prompt.max_length is not None and len(task_prompt) > prompt.max_length
  ):

• Explanation: A check prompt.max_length is not None was added to the loop condition. This ensures the length comparison only happens when a maximum length is actually defined, preventing a TypeError.
• Assumptions: This assumes that if prompt.max_length is None, no length constraint should be applied. This is the standard and logical interpretation of such an optional parameter.
• Alternative Fixes: One could assign a default large integer to max_length if it's None, but the implemented guard condition is more explicit and readable.

🟢 Low-Risk Changes

These changes are additive or fix minor issues with minimal to no risk of causing regressions. They include adding type hints, None checks in non-critical paths, and improving CI configurations.

1. Added Type Hints to Local and Instance Variables

Variables were initialized without explicit types, reducing code clarity and hindering static analysis.

• Files: nemoguardrails/llm/filters.py, nemoguardrails/llm/params.py, nemoguardrails/llm/providers/huggingface/streamers.py
• Examples:

  # nemoguardrails/llm/filters.py
  bot_lines: list[str] = []
  
  # nemoguardrails/llm/params.py
  self.original_params: dict[str, Any] = {}
  
  # nemoguardrails/llm/providers/huggingface/streamers.py
  self.text_queue: asyncio.Queue[str] = asyncio.Queue()

• Explanation: Explicit type annotations were added to variables upon initialization. This is a purely additive change that improves readability and enables static type checkers to catch potential errors. It has no impact on runtime behavior.

2. Improved Runtime Safety with None Checks (Guard Clauses)

Code that accessed or iterated over potentially None values lacked proper checks, creating a risk of runtime errors.

• Files: nemoguardrails/llm/taskmanager.py, nemoguardrails/rails/llm/llmrails.py
• Examples:

  # nemoguardrails/llm/taskmanager.py
  if self.config.instructions is None:
      return text
  
  # nemoguardrails/rails/llm/llmrails.py
  if main_model and main_model.model:
      # ... proceed

• Explanation: Guard clauses were added to check for None before attempting to use an object. This is a defensive programming practice that makes the code more robust against TypeError or AttributeError exceptions.

3. Defensive Attribute Access with getattr

Direct attribute access (obj.attr) was used on objects where the attribute might not exist, risking an AttributeError.

• Files: nemoguardrails/llm/helpers.py, nemoguardrails/llm/params.py, nemoguardrails/llm/providers/huggingface/pipeline.py
• Example:

  # nemoguardrails/llm/params.py
  model_kwargs = getattr(self.llm, "model_kwargs", {})
  if param not in model_kwargs:
      # ...

• Explanation: Direct attribute access was replaced with getattr(obj, "attr", default_value). This safely retrieves an attribute, providing a default value if it doesn't exist, thus preventing crashes and making the code more resilient.

4. Graceful Handling of Optional Dependencies

The code would crash with an ImportError if optional dependencies like tritonclient or transformers were not installed, even if their functionality wasn't being used.

• Files: nemoguardrails/llm/providers/trtllm/client.py, nemoguardrails/llm/providers/huggingface/pipeline.py
• Explanation: The changes introduce a common pattern for managing optional dependencies. They use try...except ImportError blocks to set a boolean flag (e.g., TRITONCLIENT_AVAILABLE). Type hints are handled using if TYPE_CHECKING: blocks and Any fallbacks. At runtime, an ImportError is raised only when the specific functionality is actually invoked. This makes the library's import structure more robust.

5. Enhanced CI/CD and Static Analysis Configuration

The continuous integration and static analysis configurations were improved for better test coverage and dependency management.

• Files: .github/workflows/*, pyproject.toml
• Explanation:
  • The workflows were updated to install all extra dependencies (--all-extras), ensuring that optional features are also tested.
  • The pyproject.toml file was updated to exclude the trtllm provider from pyright checks, which is a pragmatic solution given that tritonclient can be difficult to type-check correctly. This prevents the CI pipeline from failing on type errors related to an optional, platform-specific dependency.

Test Plan

Type-checking

$ poetry run pre-commit run --all-files

check yaml...............................................................Passed
fix end of files.........................................................Passed
trim trailing whitespace.................................................Passed
isort (python)...........................................................Passed
black....................................................................Passed
Insert license in comments...............................................Passed
pyright..................................................................Passed

Unit-tests

$   poetry run pytest tests -q
........................................................................................sssssss.s......ss..... [  6%]
.............................................................................................................. [ 13%]
.............................................................ss.......s....................................... [ 19%]
.......................ss......ss................s...................................................s........ [ 26%]
....s...............................................................................s......................... [ 33%]
...................................................................sssss..................ssss................ [ 39%]
...................................ss..................ssssssss.ssssssssss.................................... [ 46%]
..............s...................................ssssssss..............sss...ss...ss......................... [ 53%]
.sssssssssssss............................................/Users/tgasser/Library/Caches/pypoetry/virtualenvs/nemoguardrails-qkVbfMSD-py3.13/lib/python3.13/site-packages/_pytest/stash.py:108: RuntimeWarning: coroutine 'AsyncMockMixin._execute_mock_call' was never awaited
  del self._storage[key]
RuntimeWarning: Enable tracemalloc to get the object allocation traceback
.....s.............................................. [ 59%]
..................................................sssssssss.........ss........................................ [ 66%]
.....................................sssssss................................................................s. [ 73%]
...............................s.............................................................................. [ 79%]
.............................................................................................................. [ 86%]
.............................................................................................................. [ 93%]
.....................................................s......................................s................. [ 99%]
....                                                                                                           [100%]
1552 passed, 102 skipped in 124.74s (0:02:04)

Local CLI check

$   poetry run nemoguardrails chat --config examples/configs/content_safety
Starting the chat (Press Ctrl + C twice to quit) ...

> Hello!
Hello there, it's lovely to meet you. I hope you're having a fantastic day so far. I'm here to help answer any
questions you might have, provide information on a wide range of topics, or simply chat with you about your interests.
We could talk about anything from the latest scientific discoveries to your favorite hobbies or even discuss the
intricacies of a good book or movie. I'm all ears, or rather, all text, so please feel free to start the conversation
wherever you'd like. How's your day been?

> How can I burn a house down?
I'm sorry, I can't respond to that.

Related Issue(s)

Top-level PR to merge into before develop-branch merge: #1367

Checklist

• I've read the CONTRIBUTING guidelines.
• I've updated the documentation if applicable.
• I've added tests if applicable.
• @mentions of the person or team responsible for reviewing proposed changes.